#include <iostream>
#include <sys/fcntl.h>
#include <stdio.h>
#include <unistd.h>
#include "BPatch.h"
#include "BPatch_addressSpace.h"
#include "BPatch_process.h"
#include "BPatch_snippet.h"
#include "BPatch_binaryEdit.h"
#include "BPatch_point.h"
#include "BPatch_function.h"

// 判断是32位的机器还是64位的机器，定义SP的标准化名称
#if __SIZEOF_POINTER__ == 4
#define SP "esp"
#elif __SIZEOF_POINTER__ == 8
#define SP "rsp"
#endif

using namespace std;

bool overflowFlag;

// 查找指定寄存器的表达式
// 不知道为啥我找不到可用于寄存器查找的方法，所以用这个比较low的方法遍历查找
BPatch_registerExpr findRegister(const char *name, BPatch_addressSpace *app)
{
    vector<BPatch_register> regs;
    app->getRegisters(regs);
    int j = 0, n = regs.size();
    while (j < n && regs[j].name().compare(name) != 0)
        j++;
    return BPatch_registerExpr(regs[j]);
}

// 栈溢出回调
void overflowCallBack(BPatch_point *point, void *retValue)
{
    cout << "侦测到在途的核打击，停止程序，溢出后地址为：" << hex << retValue << endl;
    overflowFlag = true;
}

int main(int argc, char **argv)
{
    BPatch *bpatch = new BPatch;
    BPatch_process *appProc;
    vector<BPatch_function *> functions;
    vector<BPatch_point *> *entryPoints;
    vector<BPatch_point *> *retPoints;

    // LOG: init
    cout << "守卫正在初始化..." << endl;
    overflowFlag = false;

    // 输入可执行文件的名称及其输入
    appProc = bpatch->processCreate(argv[1], NULL);

    // 暂停以进行插桩
    appProc->stopExecution();

    // LOG: appProc
    cout << "appProc 加载完成" << endl;

    // 拿到地址空间
    BPatch_addressSpace *app = appProc;

    // 拿到进程镜像
    BPatch_image *appImage = app->getImage();

    // LOG: appImage
    cout << "appImage 加载完成" << endl;

    // 找到所有方法（包括Shared Lib的方法）
    appImage->getProcedures(functions);

    // 拿到方法向量的维数
    int funcs_size = functions.size();

    // 插桩开始
    app->beginInsertionSet();

    // 遍历感兴趣的方法
    for (int i = 0; i < funcs_size; i++)
    {
        // 过滤没法插入和Shared Lib里的方法
        if (!functions[i]->isInstrumentable() || functions[i]->isSharedLib())
            continue;

        // LOG：开始注入灵魂
        cout << "注入灵魂：" << functions[i]->getName() << endl;

        // 取得方法的入口和出口
        entryPoints = functions[i]->findPoint(BPatch_locEntry);
        retPoints = functions[i]->findPoint(BPatch_locExit);

        //-------------------------------------函数入口的插桩--------------------------------------
        // 为影子地址分配空间（查看汇编发现在0x700000上，在蹦床代码的地址空间0x7000f8之前）
        // 之所以用int，因为在大多数机器上偏移地址最多只需4位
        BPatch_variableExpr *shadowAddr = app->malloc(*(appImage->findType("int")));

        // 找到栈指针寄存器
        BPatch_registerExpr sp = findRegister(SP, app);

        // 拿到SP指向的内存（此时未进行任何压栈操作，也就是说仍指向call时压入的返回地址），即函数返回地址，然后赋值给影子地址
        BPatch_arithExpr fetchReturnAddr(BPatch_assign, *shadowAddr, BPatch_arithExpr(BPatch_deref, sp));

        // 插桩到函数入口
        app->insertSnippet(fetchReturnAddr, *entryPoints);

        //-------------------------------------函数出口的插桩--------------------------------------
        // 同样地，拿到SP指向的函数返回地址进行比较（此时函数栈中的本地变量都已经pop掉了，于是指回函数返回地址）
        BPatch_snippet *addr = new BPatch_arithExpr(BPatch_deref, sp);

        // 检查，得到布尔值
        BPatch_boolExpr checkReturnAddress(BPatch_ne, *shadowAddr, *addr);

        // 中断指令，捅到回调函数
        BPatch_stopThreadExpr callStop(overflowCallBack, *addr);

        // 条件跳转分支，如果检测到溢出则跳转
        BPatch_ifExpr branchIfOverflow(checkReturnAddress, callStop);

        // 插桩到函数出口
        app->insertSnippet(branchIfOverflow, *retPoints);
    }

    // 插桩完成
    app->finalizeInsertionSet(true);

    // LOG: 插桩完成，提示输入
    cout << "完成插桩，现在可以输入了（如果有的话）" << endl;

    // 被插桩程序可以继续执行了
    appProc->continueExecution();

    // 事件循环
    while (1)
    {
        // wait一把
        if (bpatch->waitForStatusChange())
        {
            // 判断感兴趣的事件
            if (appProc->isTerminated())
            {
                appProc->detach(1);
                exit(0);
            }
            else if (appProc->isStopped())
            {
                cout << "thread stop..." << endl;
            }
            else if (overflowFlag)
            {
                appProc->terminateExecution();
                appProc->detach(1);
                exit(0);
            }
        }
    }
}